# -*- coding: UTF-8 -*-
import time
import shutil
import socket
import select
import torch
from torchvision import datasets, models, transforms
from PIL import Image
from PIL import ImageFile
import os
# 导入Python标准库中的Process模块

from threading import Thread


def utf8(text):
    return str(text, 'utf-8')


def gbk(text):
    return text.decode("gbk", "ignore")


def threads_predict(*pics):
    # 发送的报文
    message = []
    for file in pics:
        img = Image.open(cache_path + file).convert('RGB')
        img = transform(img)
        img = torch.unsqueeze(img, dim=0)
        with torch.no_grad():
            outputs = net(img)
            # 得到预测分类结果
            predict = torch.max(outputs, dim=1)[1].data.numpy()
            # 如果概率低于0.7, 认为是其它类
            maxPercent = float(torch.softmax(outputs, dim=1)[0].numpy()[predict])
            if maxPercent <= 0.5 or classes[int(predict)] == 'document':
                predict = 8
            # 填充报文
            message.append(str({"photoName": file, "classify": classes[int(predict)]}))
            # 移动文件
            if classes[int(predict)] == "people":
                shutil.move(cache_path + file, aim_path + classes[int(predict)] + '/cache/')
            else:
                shutil.move(cache_path + file, aim_path + classes[int(predict)] + '/')
    # 报文处理
    paper = "[" + ",".join(message) + "]"
    event.send(paper.encode())


ImageFile.LOAD_TRUNCATED_IMAGES = True
# 线程池
threads = []


transform = transforms.Compose([
        # 按照比例把图像最小的一个边长放缩到256，另一边按照相同比例放缩。
        transforms.Resize(256),
        # 在图片的中间区域进行裁剪 224*224
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])

# 类别
classes = ('animal', 'building', 'comic', 'document', 'food', 'people', 'traffic', 'view', 'other')

# 加载模型
net = torch.load('./eightType.pkl')
print("模型加载完成")

# 图片缓存的固定路径
fixed_path = "D:/JAVA-competition/test/wisdom/"
cache_path = ""

# select-socket通信建立
sock = socket.socket()
sock.bind(('127.0.0.1', 3333))
sock.setblocking(False)
sock.listen(10)

inputs = [sock, ]

while True:
    r_list, w_list, e_list = select.select(inputs, [], [], 1)
    for event in r_list:
        if event == sock:
            print("新的客户端连接")
            new_sock, addr = event.accept()
            inputs.append(new_sock)
        else:
            data = event.recv(1024)
            if data:
                print("命令:", utf8(data))
                # 临时文件夹路径
                cache_path = fixed_path + utf8(data) + '/' + 'cache/'
                # 目标文件夹路径
                aim_path = fixed_path + utf8(data) + '/'
                # 读取
                start = time.time()
                if os.path.exists(cache_path):
                    files = os.listdir(cache_path)
                    print("一共有" + str(files.__len__()) + "张图片")
                    if files.__len__() != 0:
                        threads_nums = 1
                        # 根据图片张数决定线程个数
                        if files.__len__() > 20:
                            threads_nums = int(files.__len__() / 20) + 1
                        for i in range(threads_nums):
                            if files.__len__() - i*20 < 20:
                                thread = Thread(target=threads_predict, args=files[i * 20:files.__len__()])
                            else:
                                thread = Thread(target=threads_predict, args=files[i * 20:(i + 1) * 20])
                            threads.append(thread)
                            thread.setDaemon(True)
                        for i in range(threads_nums):
                            threads[i].start()
                        for i in range(threads_nums):
                            threads[i].join()
                        threads = []
                        event.send("over".encode())
                        print("秒数为：" + str(time.time() - start))
                else:
                    print("缓存路径缺失")
            else:
                print("客户端断开连接")
                inputs.remove(event)
